/*
	Copyright (C) 2008 Christopher E. Miller
	
	This software is provided 'as-is', without any express or implied
	warranty.  In no event will the authors be held liable for any damages
	arising from the use of this software.
	
	Permission is granted to anyone to use this software for any purpose,
	including commercial applications, and to alter it and redistribute it
	freely, subject to the following restrictions:
	
	1. The origin of this software must not be misrepresented; you must not
	   claim that you wrote the original software. If you use this software
	   in a product, an acknowledgment in the product documentation would be
	   appreciated but is not required.
	2. Altered source versions must be plainly marked as such, and must not be
	   misrepresented as being the original software.
	3. This notice may not be removed or altered from any source distribution.
*/


/*  gcc -ldl -rdynamic -o ircly -g ircly.c */


#include <stdlib.h>
#include <string.h>
#include <stdio.h>

#include <sys/types.h>
#include <sys/socket.h>
#include <sys/time.h>
#include <unistd.h>
#include <signal.h>
#include <dlfcn.h>
#include <fcntl.h>
#include <sys/stat.h>
#include <errno.h>


void* server_malloc(size_t len)
{
	return malloc(len);
}


void* server_realloc(void* p, size_t len)
{
	return realloc(p, len);
}


void server_free(void* p)
{
	free(p);
}


unsigned sendbuf_size = 1024 * 12;
unsigned recvbuf_size = 1024 * 12;


int attached = 2; /* 0 = not, 1 = yes, 2 = reattaching service, 3 is detach then reattach (2 steps). */


#define SFLAG_IRC 1
#define SFLAG_BUFFERED 2
#define SFLAG_ERROR 0x40000000
#define SFLAG_CLOSED 0x80000000


/* Note! callback can be NULL, and is NOT VALID if the service is not attached!
   callback should return 0 unless otherwise instructed.
   callback id is 1 for read, cinfo is pointer to int with recv value.
*/
typedef struct _SockInfo
{
	int socket;
	unsigned long flags;
	char* name;
	void* data;
	size_t datalen;
	int (*callback)(struct _SockInfo*, unsigned id, void* cinfo);
	char* sendbuf;
	size_t sendlen;
	char* nextsend;
	char* recvbuf;
	size_t recvlen;
	char* nextrecv;
}SockInfo;


SockInfo** sinfos = NULL;
size_t nsinfos = 0;


size_t socket_get_all(void (*callback)(SockInfo*, void* extra), void* extra)
{
	size_t iw;
	if(callback)
	{
		for(iw = 0; iw != nsinfos; iw++)
		{
			callback((SockInfo*)sinfos[iw], extra);
		}
	}
	return nsinfos;
}


size_t socket_get_array(SockInfo*** pinfos)
{
	if(pinfos)
		*pinfos = (SockInfo**)sinfos;
	return nsinfos;
}


SockInfo* socket_info(int socket)
{
	size_t iw;
	for(iw = 0; iw != nsinfos; iw++)
	{
		if(socket == sinfos[iw]->socket)
			return sinfos[iw];
	}
	return NULL;
}


static void _sinfo_dealloc(SockInfo* sinfo)
{
	/* It's allocated all-in-one now.
	if(sinfo->name)
		free(sinfo->name);
	if(sinfo->data)
		free(sinfo->data);
	if(sinfo->sendbuf)
		free(sinfo->sendbuf);
	if(sinfo->recvbuf)
		free(sinfo->recvbuf);
	*/
	free(sinfo);
}


/* Set flags to 1 for IRC socket, or 0 for any old socket.
	Setting flags to 2 will allow a non-IRC socket to be buffered.
*/
int socket_register(int socket, unsigned long flags, char* name,
	int (*callback)(SockInfo*, unsigned id, void* cinfo))
{
	SockInfo* sinfo;
	SockInfo** newsinfos;
	char* cs;
	size_t n, n2;
	
	sinfo = socket_info(socket);
	if(sinfo)
		return 2;
	
	n = sizeof(SockInfo);
	
	if(!name)
		name = "";
	n2 = strlen(name) + 1;
	n += n2;
	
	if(flags & (1 | 2))
	{
		n += sendbuf_size;
		n += recvbuf_size;
	}
	
	cs = malloc(n);
	if(!cs)
		return 0;
	
	sinfo = (SockInfo*)cs;
	cs += sizeof(SockInfo);
	memset(sinfo, 0, sizeof(SockInfo)); /* Zero! */
	
	sinfo->name = cs;
	cs += n2;
	memcpy(sinfo->name, name, n2);
	
	if(flags & (1 | 2))
	{
		sinfo->sendbuf = sinfo->nextsend = cs;
		cs += sendbuf_size;
		sinfo->recvbuf = sinfo->nextrecv = cs;
		/*cs += recvbuf_size;*/
	}
	
	sinfo->socket = socket;
	sinfo->flags = flags;
	sinfo->callback = callback;
	
	newsinfos = realloc(sinfos, (nsinfos + 1) * sizeof(SockInfo*));
	if(!newsinfos)
		goto sfail;
	newsinfos[nsinfos] = sinfo;
	sinfos = newsinfos;
	nsinfos++;
	
	return 1;
	
	sfail:
	_sinfo_dealloc(sinfo);
	return 0;
}


int socket_set_callback(int socket,
	int (*callback)(SockInfo*, unsigned id, void* cinfo))
{
	SockInfo* sinfo;
	sinfo = socket_info(socket);
	if(!sinfo)
		return 0;
	sinfo->callback = callback;
	return 1;
}


static int _bufrecv(SockInfo* sinfo)
{
	int i;
	
	if(!sinfo->recvlen)
		sinfo->nextrecv = sinfo->recvbuf;
	
	if(sinfo->nextrecv != sinfo->recvbuf)
	{
		if(recvbuf_size - sinfo->recvlen - (sinfo->nextrecv - sinfo->recvbuf) < 512)
		{
			memmove(sinfo->recvbuf, sinfo->nextrecv, sinfo->recvlen);
			sinfo->nextrecv = sinfo->recvbuf;
		}
	}
	else /* Next (user)recv at beginning. */
	{
		if(sinfo->recvlen >= recvbuf_size)
		{
			return -1; /* Out of buffer space! */
		}
	}
	
	i = recv(sinfo->socket,
		sinfo->nextrecv + sinfo->recvlen,
		recvbuf_size - ((sinfo->nextrecv + sinfo->recvlen) - sinfo->recvbuf),
		0);
	if(!i)
	{
		sinfo->flags |= SFLAG_CLOSED;
		return 0;
	}
	if(i < 0)
	{
		sinfo->flags |= SFLAG_ERROR;
		return i;
	}
	sinfo->recvlen += i;
	return i;
}


static int _bufsend(SockInfo* sinfo)
{
	int i;
	
	i = send(sinfo->socket, sinfo->nextsend, sinfo->sendlen, 0);
	
	if(i > 0)
	{
		sinfo->sendlen -= i;
		sinfo->nextsend += i;
	}
	
	return i;
}


/* Negative = error, 0 = closed, 1 = connected.
	Note that the state is only affected by calling these socket functions.
*/
int socket_state(int socket)
{
	SockInfo* sinfo;
	sinfo = socket_info(socket);
	if(sinfo)
	{
		if(sinfo->flags & SFLAG_CLOSED)
			return 0;
		if(sinfo->flags & SFLAG_ERROR)
			return -1;
		return 1;
	}
	return -2;
}


/* Not all socket flags are supported. MSG_PEEK is supported for buffered.
	len is updated with the recv result; -1 on error, 0 on close.
	NULL return is non-socket error (probably memory error).
*/
char* socket_recv(int socket, int* len, unsigned flags)
{
	SockInfo* sinfo;
	void* p;
	char* cs;
	int i;
	
	if(!len)
		return NULL;
	if(*len <= 0)
	{
		xbad:
		*len = 0;
		return NULL;
	}
	
	sinfo = socket_info(socket);
	if(!sinfo)
		goto xbad;
	
	if(sinfo->flags & (1 | 2))
	{
		if(!sinfo->recvlen)
		{
			if(_bufrecv(sinfo) < 0)
			{
				*len = -1;
				return sinfo->nextrecv;
			}
		}
		
		if(*len > sinfo->recvlen)
			*len = sinfo->recvlen;
		if(flags & MSG_PEEK)
			return sinfo->nextrecv;
		sinfo->recvlen -= *len;
		cs = sinfo->nextrecv;
		sinfo->nextrecv += *len;
		return cs;
	}
	else
	{
		if(*len > 4 * 1024)
			*len = 4 * 1024;
		
		if(sinfo->recvbuf)
		{
			if(*len > sinfo->recvlen)
			{
				/* Note: this causes unwanted copying.
				p = realloc(sinfo->recvbuf, *len);
				if(!p)
					goto xbad;
				sinfo->recvlen = *len;
				*/
				free(sinfo->recvbuf);
				sinfo->recvlen = 0;
				goto need_recvbuf_alloc;
			}
		}
		else
		{
			need_recvbuf_alloc:
			sinfo->recvbuf = malloc(*len);
			if(!sinfo->recvbuf)
				goto xbad;
			sinfo->recvlen = *len;
		}
		
		*len = recv(socket, sinfo->recvbuf, *len, flags);
		if(!*len)
			sinfo->flags |= SFLAG_CLOSED;
		if(*len < 0)
			sinfo->flags |= SFLAG_ERROR;
		return sinfo->recvbuf;
	}
}


int socket_send(int socket, const void* buf, int len, unsigned flags)
{
	SockInfo* sinfo;
	int i;
	
	sinfo = socket_info(socket);
	if(!sinfo)
		return -1;
	
	if(!len)
		return 0;
	
	if(sinfo->flags & (1 | 2))
	{
		i = 0; /* Important! */
		if(!sinfo->sendlen)
		{
			i = send(socket, buf, len, flags);
			if(i < 0)
			{
				/* Don't treat this as a real error, because it might not actually be time to send yet,
				   and that would be this code's fault... so just buffer it and let it pass/fail later.
				*/
				/*return i;*/
				i = 0;
			}
			if(i == len)
				return i;
			buf += i;
			len -= i;
		}
		/* i must be preserved! */
		if(len > sendbuf_size - (sinfo->sendlen - (sinfo->nextsend - sinfo->sendbuf)))
		{
			if(sinfo->sendbuf != sinfo->nextsend)
			{
				memmove(sinfo->sendbuf, sinfo->nextsend, sinfo->sendlen);
				sinfo->nextsend = sinfo->sendbuf;
			}
			
			/* If it still doesn't fit, the buffer is full... so just truncate it and report. */
			if(len > sendbuf_size - sinfo->sendlen)
				len = sendbuf_size - sinfo->sendlen;
		}
		memcpy(sinfo->nextsend + sinfo->sendlen, buf, len);
		sinfo->sendlen += len;
		return i + len;
	}
	else
	{
		i = send(socket, buf, len, flags);
		if(i < 0)
			sinfo->flags |= SFLAG_ERROR;
		return i;
	}
}


void socket_close(int socket)
{
	SockInfo* sinfo;
	
	close(socket);
	
	sinfo = socket_info(socket);
	if(!sinfo)
		return;
	
	sinfo->flags |= SFLAG_CLOSED;
}


void socket_unregister(int socket)
{
	SockInfo* sinfo;
	size_t iw;
	
	for(iw = 0; iw != nsinfos; iw++)
	{
		if(socket == sinfos[iw]->socket)
		{
			sinfo = sinfos[iw];
			goto got_sinfo;
		}
	}
	return;
	got_sinfo: ;
	
	_sinfo_dealloc(sinfo);
	
	/*if(iw != nsinfos - 1)*/
		sinfos[iw] = sinfos[nsinfos - 1];
	nsinfos--;
	sinfos = realloc(sinfos, nsinfos * sizeof(SockInfo*));
	if(!sinfos)
		nsinfos = 0;
}


/* Preserves original contents if data previously existed. */
void* socket_alloc_data(int socket, size_t datalen)
{
	SockInfo* sinfo;
	void* p;
	
	sinfo = socket_info(socket);
	if(!sinfo)
		return NULL;
	
	if(!datalen)
	{
		if(sinfo->data)
		{
			free(sinfo->data);
			sinfo->data = NULL;
			sinfo->datalen = 0;
		}
		return NULL;
	}
	
	p = realloc(sinfo->data, datalen);
	if(!p)
		return NULL;
	sinfo->data = p;
	sinfo->datalen = datalen;
	return p;
}


/* pdatalen can be NULL */
void* socket_get_data(int socket, size_t* pdatalen)
{
	SockInfo* sinfo;
	
	sinfo = socket_info(socket);
	if(!sinfo)
	{
		if(pdatalen)
			*pdatalen = 0;
		return NULL;
	}
	
	if(pdatalen)
		*pdatalen = sinfo->datalen;
	return sinfo->data;
}


int socket_select(int nfds, fd_set* reads, fd_set* writes, fd_set* errors, struct timeval* ptv)
{
	fd_set myreads, mywrites;
	struct timeval tv;
	size_t iw;
	char* cs;
	int mynfds = -1;
	int result;
	int i;
	
	if(!reads)
	{
		FD_ZERO(&myreads);
		reads = &myreads;
	}
	if(!writes)
	{
		FD_ZERO(&mywrites);
		writes = &mywrites;
	}
	
	/* Timeout... by default can't wait too long because need to check stuff. */
	if(!ptv)
	{
		tv.tv_sec = 0;
		tv.tv_usec = 500000; /* 0.5 seconds = 500 000 microseconds */
		ptv = &tv;
	}
	else
	{
		if(tv.tv_sec)
		{
			tv.tv_sec = 0;
			tv.tv_usec = 500000;
		}
		else if(tv.tv_usec > 500000)
		{
			tv.tv_usec = 500000;
		}
	}
	
	for(iw = 0; iw != nsinfos; iw++)
	{
		if(sinfos[iw]->flags & (1 | 2))
		{
			if(!(sinfos[iw]->flags & (SFLAG_ERROR | SFLAG_CLOSED)))
			{
				if(sinfos[iw]->sendlen)
				{
					FD_SET(sinfos[iw]->socket, writes);
					if(sinfos[iw]->socket > mynfds)
						mynfds = sinfos[iw]->socket;
				}
				if(sinfos[iw]->recvlen < recvbuf_size) /* Only if there's room! */
				{
					FD_SET(sinfos[iw]->socket, reads);
					if(sinfos[iw]->socket > mynfds)
						mynfds = sinfos[iw]->socket;
				}
			}
		}
	}
		
	mynfds++;
	if(mynfds > nfds)
		nfds = mynfds;
	
	result = select(nfds, reads, writes, errors, ptv);
	
	if(result > 0)
	{
		for(iw = 0; iw != nsinfos; iw++)
		{
			if(sinfos[iw]->flags & (1 | 2))
			{
				if(sinfos[iw]->sendlen)
				{
					if(FD_ISSET(sinfos[iw]->socket, writes))
					{
						_bufsend(sinfos[iw]);
					}
				}
				if(FD_ISSET(sinfos[iw]->socket, reads))
				{
					i = _bufrecv(sinfos[iw]);
					if(1 == attached && sinfos[iw]->callback)
					{
						sinfos[iw]->callback(sinfos[iw], 1, &i);
					}
				}
			}
		}
	}
	
	return result;
}


/* Detach service. */
static void _sigusr1(int signum)
{
	attached = 0;
}


/* Attach service. */
static void _sigusr2(int signum)
{
	attached = 2;
}


/* Reattach: detach and attach service. */
static void _sighup(int signum)
{
	attached = 3;
}


/* Returns 0 on failure. */
static int daemonize(const char *parentmsg)
{
	int i;
	if(1 != getppid())
	{
		if(parentmsg)
			printf("%s", parentmsg);
		
		i = fork();
		if (i < 0)
			return 0; /* Unable to fork. */
		if(i > 0)
			exit(0);
		
		setsid(); /* New process group. */
		
		close(0);
		close(1);
		close(2);
		
		i = open("/dev/null", O_RDWR); /* 0 - in */
		i = open("out.log", O_CREAT | O_WRONLY, S_IRUSR | S_IWUSR | S_IRGRP); /* 1 - out */
		if(i < 0)
			i = open("/dev/null", O_RDWR);
		else
			lseek(i, 0, SEEK_END); /* Append. */
		dup(i); /* 2 - err */
		
		/* Ignore a few signals. */
		signal(SIGCHLD, SIG_IGN);
		signal(SIGTSTP, SIG_IGN);
		signal(SIGTTOU, SIG_IGN);
		signal(SIGTTIN, SIG_IGN);
	}
	return 1;
}


int main(int argc, char** argv)
{
	void* dl = NULL;
	int (*dlrunfunc)(int) = NULL;
	int i;
	
	if(argc > 1 && '+' == argv[1][0])
	{
		printf("Not turning into daemon\n");
	}
	else
	{
		if(!daemonize("Turning into daemon...\n"))
		{
			fprintf(stderr, "Unable to turn into daemon\n");
			return EXIT_FAILURE;
		}
	}
	
	signal(SIGPIPE, SIG_IGN); /* ! */
	
	signal(SIGUSR1, _sigusr1);
	signal(SIGUSR2, _sigusr2);
	signal(SIGHUP, _sighup);
	
	printf("%s startup\n", argv[0]);
	
	/* Sorry this code is a bit spaghettilike. */
	for(;;)
	{
		switch(attached)
		{
			case 0: /* Detached. */
				detached:
				if(dl)
				{
					if(dlrunfunc)
						dlrunfunc(2);
					dlclose(dl);
					dl = NULL;
					dlrunfunc = NULL;
					printf("Service detached\n");
					/* Do NOT set 'attached' to 0 here! */
					break;
				}
				
				defrun:
				socket_select(0, NULL, NULL, NULL, NULL);
				break;
			
			case 1: /* Attached. */
				if(dl && dlrunfunc)
				{
					i = dlrunfunc(0);
					if(-42 == i)
					{
						attached = 0;
						goto detached;
					}
					if(i)
						break;
				}
				goto defrun;
			
			case 2: /* Attach! */
				if(dl)
					goto detached; /* Detach but stay at 2. */
				
				dl = dlopen("./irclyserv.so", RTLD_LAZY);
				if(!dl)
				{
					fprintf(stderr, "%s: unable to load service: %s\n", argv[0], dlerror());
					attached = 0;
					break;
				}
				
				dlrunfunc = dlsym(dl, "service_run");
				if(!dlrunfunc)
				{
					fprintf(stderr, "%s: unable to load service: missing service_run\n", argv[0]);
					attached = 0;
					goto detached;
				}
				
				if(!dlrunfunc(1))
				{
					fprintf(stderr, "%s: unable to load service: load procedure returned false\n", argv[0]);
					attached = 0;
					dlrunfunc = NULL; /* Don't "unload" if didn't load. */
					goto detached;
				}
				
				attached = 1;
				
				printf("Service attached\n");
				break;
			
			case 3: /* Reattaching.. */
				attached = 2;
				goto detached;
			
			default: /* ..? */
				return 1;
		}
	}
	
	return 0;
}

